from transformers import GPT2LMHeadModel, GPT2Tokenizer

class LyricsGenerator:
    def __init__(self):
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
        self.model = GPT2LMHeadModel.from_pretrained("gpt2")

    def generate(self, prompt, max_length=100):
        inputs = self.tokenizer(prompt, return_tensors="pt")
        return self.tokenizer.decode(
            self.model.generate(**inputs, max_length=max_length)[0]
        )